import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from statsmodels.tsa.seasonal import seasonal_decompose
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from statsmodels.tsa.statespace.sarimax import SARIMAX
from statsmodels.tsa.stattools import adfuller
from tqdm import tqdm

# 读取CSV文件
file_path = 'Demodata.csv'
df = pd.read_csv(file_path)

# 将Date列转换为日期时间格式
df['Date'] = pd.to_datetime(df['Date'])

# 设置Date列为索引
df.set_index('Date', inplace=True)

# 提取h1列的数据
h1_column = df['h1']

# 1. 数据可视化
plt.figure(figsize=(14, 8))

# 绘制原始数据
plt.subplot(2, 2, 1)
plt.plot(h1_column, label='h1', color='blue', marker='o')
plt.title('Trend of h1 Column')
plt.xlabel('Date')
plt.ylabel('Values')
plt.legend()
plt.grid(True)

# 2. 移动平均线
rolling_mean = h1_column.rolling(window=12).mean()
plt.subplot(2, 2, 1)
plt.plot(rolling_mean, label='12-Month Rolling Mean', color='red')
plt.legend()

# 3. 统计分析
mean_value = h1_column.mean()
median_value = h1_column.median()
std_dev = h1_column.std()

print(f"\nh1 Column Statistics:")
print(f"Mean: {mean_value}")
print(f"Median: {median_value}")
print(f"Standard Deviation: {std_dev}")

# 4. 季节性和周期性分析
decomposition = seasonal_decompose(h1_column, model='additive', period=12)
trend = decomposition.trend
seasonal = decomposition.seasonal
residual = decomposition.resid

# 绘制分解结果
plt.subplot(2, 2, 2)
plt.plot(trend, label='Trend', color='green')
plt.title('Trend Component')
plt.xlabel('Date')
plt.ylabel('Values')
plt.legend()
plt.grid(True)

plt.subplot(2, 2, 3)
plt.plot(seasonal, label='Seasonal', color='orange')
plt.title('Seasonal Component')
plt.xlabel('Date')
plt.ylabel('Values')
plt.legend()
plt.grid(True)

plt.subplot(2, 2, 4)
plt.plot(residual, label='Residual', color='purple')
plt.title('Residual Component')
plt.xlabel('Date')
plt.ylabel('Values')
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()

# 5. 自相关分析
plt.figure(figsize=(14, 6))

plt.subplot(2, 1, 1)
plot_acf(h1_column, lags=30, ax=plt.gca())
plt.title('Autocorrelation Function (ACF) of h1 Column')

plt.subplot(2, 1, 2)
plot_pacf(h1_column, lags=30, ax=plt.gca())
plt.title('Partial Autocorrelation Function (PACF) of h1 Column')

plt.tight_layout()
plt.show()

# 6. 检查平稳性
def adf_test(series):
    result = adfuller(series)
    print('ADF Statistic:', result[0])
    print('p-value:', result[1])
    print('Critical Values:')
    for key, value in result[4].items():
        print(f'\t{key}: {value}')
    return result

# 检查原始数据的平稳性
adf_result = adf_test(h1_column)

# 7. 差分处理（如果需要）
if adf_result[1] > 0.05:
    h1_diff = h1_column.diff().dropna()
    adf_result_diff = adf_test(h1_diff)
else:
    h1_diff = h1_column
    adf_result_diff = adf_result

# 8. 自动选择ARIMA参数
# 使用SARIMAX模型进行网格搜索
p_values = range(0, 3)
d_values = range(0, 2)
q_values = range(0, 3)

best_aic = float("inf")
best_order = None

# 使用tqdm显示进度条
for p in tqdm(p_values, desc="Searching p"):
    for d in tqdm(d_values, desc="Searching d", leave=False):
        for q in tqdm(q_values, desc="Searching q", leave=False):
            try:
                model = SARIMAX(h1_diff, order=(p, d, q))
                model_fit = model.fit(disp=False)  # 设置disp=False以不显示具体的计算过程
                current_aic = model_fit.aic
                if current_aic < best_aic:
                    best_aic = current_aic
                    best_order = (p, d, q)
            except:
                continue

print(f"\nBest ARIMA parameters: p={best_order[0]}, d={best_order[1]}, q={best_order[2]}")

# 9. ARIMA 模型训练
model = SARIMAX(h1_diff, order=best_order)
model_fit = model.fit(disp=False)  # 设置disp=False以不显示具体的计算过程

print(model_fit.summary())

# 10. 预测下一个数据点
forecast = model_fit.forecast(steps=1)
next_value = forecast.iloc[0]  # 使用iloc访问预测值

# 如果进行了差分处理，需要反差分
if best_order[1] > 0:
    last_value = h1_diff.iloc[-1]
    next_value = last_value + next_value

# 对预测结果进行四舍五入取整
next_value = round(next_value)

print(f"\nPredicted next value for h1: {next_value}")

# 11. 绘制局部趋势和预测结果
plt.figure(figsize=(10, 6))
plt.plot(h1_column.index, h1_column, label='Original', color='blue', marker='o')
plt.plot(pd.date_range(start=h1_column.index[-1], periods=2, freq='D')[1], next_value, label='Predicted Next Value', color='red', marker='x', markersize=10)
plt.title('Local Trend and Prediction of h1 Column')
plt.xlabel('Date')
plt.ylabel('Values')
plt.legend()
plt.grid(True)
plt.show()